{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4/4 [==============================] - 0s 1ms/step - loss: 0.5337\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\n",
      "WARNING:tensorflow:From D:\\Anaconda3\\envs\\keras\\lib\\site-packages\\tensorflow\\python\\ops\\resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "INFO:tensorflow:Assets written to: my_model\\assets\n",
      "4/4 [==============================] - 0s 750us/step - loss: 0.4593\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\n"
     ]
    },
    {
     "data": {
      "text/plain": "<tensorflow.python.keras.callbacks.History at 0x1d8135c7188>"
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_model():\n",
    "    inputs = keras.Input(shape=(32,))\n",
    "    outputs = keras.layers.Dense(1)(inputs)\n",
    "    model = keras.Model(inputs,outputs)\n",
    "    model.compile(optimizer='adam',loss='mean_squared_error')\n",
    "    return model\n",
    "\n",
    "model = get_model()\n",
    "\n",
    "test_input = np.random.random((128,32))\n",
    "test_target = np.random.random((128,1))\n",
    "model.fit(test_input,test_target)\n",
    "\n",
    "model.save(\"my_model\")\n",
    "\n",
    "reconsstructed_model = keras.models.load_model('my_model')\n",
    "\n",
    "np.testing.assert_allclose(\n",
    "    model.predict(test_input),reconsstructed_model.predict(test_input)\n",
    ")\n",
    "\n",
    "\n",
    "reconsstructed_model.fit(test_input,test_target)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: my_model\\assets\n",
      "WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.\n",
      "Original model: <__main__.CustomModel object at 0x000001D812D3CE88>\n",
      "Loaded model: <tensorflow.python.keras.saving.saved_model.load.CustomModel object at 0x000001D81CA34F48>\n"
     ]
    }
   ],
   "source": [
    "class CustomModel(keras.Model):\n",
    "    def __init__(self,hidden_units):\n",
    "        super(CustomModel,self).__init__()\n",
    "        self.dense_layers = [keras.layers.Dense(u) for u in hidden_units]\n",
    "\n",
    "    def call(self,inputs):\n",
    "        x = inputs\n",
    "        for layer in self.dense_layers:\n",
    "            x = layer(x)\n",
    "        return x\n",
    "\n",
    "model = CustomModel([16,16,10])\n",
    "input_arr = tf.random.uniform((1,5))\n",
    "outputs = model(input_arr)\n",
    "model.save(\"my_model\")\n",
    "\n",
    "del CustomModel\n",
    "\n",
    "loaded = keras.models.load_model('my_model')\n",
    "np.testing.assert_allclose(loaded(input_arr),outputs)\n",
    "\n",
    "print(\"Original model:\",model)\n",
    "print(\"Loaded model:\",loaded)\n",
    "\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4/4 [==============================] - 0s 1ms/step - loss: 0.7072\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\n",
      "4/4 [==============================] - 0s 1000us/step - loss: 0.6194\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\n"
     ]
    },
    {
     "data": {
      "text/plain": "<tensorflow.python.keras.callbacks.History at 0x1d813a5a708>"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = get_model()\n",
    "\n",
    "test_input = np.random.random((128,32))\n",
    "test_target = np.random.random((128,1))\n",
    "model.fit(test_input,test_target)\n",
    "\n",
    "model.save(\"my_h5_model.h5\")\n",
    "\n",
    "reconsstructed_model = keras.models.load_model('my_h5_model.h5')\n",
    "\n",
    "np.testing.assert_allclose(\n",
    "    model.predict(test_input),\n",
    "    reconsstructed_model.predict(test_input)\n",
    ")\n",
    "\n",
    "reconsstructed_model.fit(test_input,test_target)\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "layer = keras.layers.Dense(3,activation='relu')\n",
    "layer_config = layer.get_config()\n",
    "new_layer = keras.layers.Dense.from_config(layer_config)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [
    "model = keras.Sequential([\n",
    "    keras.Input((32,)),\n",
    "    keras.layers.Dense(1)\n",
    "])\n",
    "config = model.get_config()\n",
    "new_model = keras.Sequential.from_config(config)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [],
   "source": [
    "inputs = keras.Input((32,))\n",
    "outputs = keras.layers.Dense(1)(inputs)\n",
    "model = keras.Model(inputs,outputs)\n",
    "config = model.get_config()\n",
    "new_model = keras.Model.from_config(config)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [],
   "source": [
    "model = keras.Sequential([\n",
    "    keras.Input((32,)),\n",
    "    keras.layers.Dense(1)\n",
    "])\n",
    "json_config = model.to_json()\n",
    "new_model = keras.models.model_from_json(json_config)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: my_model\\assets\n"
     ]
    }
   ],
   "source": [
    "model.save('my_model')\n",
    "tensorflow_graph = tf.saved_model.load(\"my_model\")\n",
    "x = np.random.uniform(size=(4,32)).astype(np.float32)\n",
    "predicted = tensorflow_graph(x).numpy()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [],
   "source": [
    "class CustomLayer(keras.layers.Layer):\n",
    "    def __init__(self,a):\n",
    "        self.var = tf.Variable(a,name='var_a')\n",
    "\n",
    "    def call(self,inputs,training=False):\n",
    "        if training:\n",
    "            return inputs * self.var\n",
    "        else:\n",
    "            return inputs\n",
    "\n",
    "    def get_config(self):\n",
    "        return {\"a\":self.var.numpy()}\n",
    "\n",
    "    @classmethod\n",
    "    def from_config(cls,config):\n",
    "        return cls(**config)\n",
    "\n",
    "layer = CustomLayer(5)\n",
    "layer.var.assign(2)\n",
    "\n",
    "serialized_layer = keras.layers.serialize(layer)\n",
    "new_layer = keras.layers.deserialize(\n",
    "    serialized_layer,custom_objects={\"CustomLayer\":CustomLayer}\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "outputs": [],
   "source": [
    "class CustomLayer(keras.layers.Layer):\n",
    "    def __init__(self,units=32,**kwargs):\n",
    "        super(CustomLayer,self).__init__(**kwargs)\n",
    "        self.units = units\n",
    "\n",
    "    def build(self,input_shape):\n",
    "        self.w = self.add_weight(\n",
    "            shape=(input_shape[-1],self.units),\n",
    "            initializer='random_normal',\n",
    "            trainable=True,\n",
    "        )\n",
    "        self.b = self.add_weight(\n",
    "            shape=(self.units,),\n",
    "            initializer='random_normal',\n",
    "            trainable=True\n",
    "        )\n",
    "\n",
    "    def call(self,inputs):\n",
    "        return tf.matmul(inputs,self.w) + self.b\n",
    "\n",
    "    def get_config(self):\n",
    "        config = super(CustomLayer,self).get_config()\n",
    "        config.update({\"units\":self.units})\n",
    "        return config\n",
    "\n",
    "def custom_activation(x):\n",
    "    return tf.nn.tanh(x) ** 2\n",
    "\n",
    "inputs = keras.Input((32,))\n",
    "x = CustomLayer(32)(inputs)\n",
    "outputs = keras.layers.Activation(custom_activation)(x)\n",
    "model = keras.Model(inputs,outputs)\n",
    "\n",
    "config = model.get_config()\n",
    "\n",
    "custom_objects = {\"CustomLayer\":CustomLayer,\"custom_activation\":custom_activation}\n",
    "with keras.utils.custom_object_scope(custom_objects):\n",
    "    new_model = keras.Model.from_config(config)\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "outputs": [],
   "source": [
    "def create_layer():\n",
    "    layer = keras.layers.Dense(64,activation='relu',name='dense_2')\n",
    "    layer.build((None,784))\n",
    "    return layer\n",
    "\n",
    "layer_1 = create_layer()\n",
    "layer_2 = create_layer()\n",
    "\n",
    "layer_2.set_weights(layer_1.get_weights())\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "outputs": [],
   "source": [
    "inputs = keras.Input(shape=(784,),name='digits')\n",
    "x = keras.layers.Dense(64,activation='relu',name='dense_1')(inputs)\n",
    "x = keras.layers.Dense(64,activation='relu',name='dense_2')(x)\n",
    "outputs = keras.layers.Dense(10,name='predictions')(x)\n",
    "functional_model = keras.Model(inputs=inputs,outputs=outputs,name='3_layer_mlp')\n",
    "\n",
    "class SubclassedModel(keras.Model):\n",
    "    def __init__(self,output_dim,name=None):\n",
    "        super(SubclassedModel,self).__init__(name=name)\n",
    "        self.output_dim = output_dim\n",
    "        self.dense_1 = keras.layers.Dense(64,activation='relu',name='dense_1')\n",
    "        self.dense_2 = keras.layers.Dense(64,activation='relu',name='dense_2')\n",
    "        self.dense_3 = keras.layers.Dense(output_dim,name='predictions')\n",
    "\n",
    "    def call(self,inputs):\n",
    "        x = self.dense_1(inputs)\n",
    "        x = self.dense_2(x)\n",
    "        x = self.dense_3(x)\n",
    "        return x\n",
    "\n",
    "    def get_config(self):\n",
    "        return {\"output_dim\":self.output_dim,'name':self.name}\n",
    "\n",
    "subclassed_model = SubclassedModel(10)\n",
    "subclassed_model(tf.ones((1,784)))\n",
    "subclassed_model.set_weights(functional_model.get_weights())\n",
    "\n",
    "assert  len(functional_model.weights) == len(subclassed_model.weights)\n",
    "for a,b in zip(functional_model.weights,subclassed_model.weights):\n",
    "    np.testing.assert_allclose(a.numpy(),b.numpy())\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "outputs": [
    {
     "data": {
      "text/plain": "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x1d8206eb388>"
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sequential_model = keras.Sequential([\n",
    "    keras.Input(shape=(784,),name='digits'),\n",
    "    keras.layers.Dense(64,activation='relu',name='dense_1'),\n",
    "    keras.layers.Dense(64,activation='relu',name='dense_2'),\n",
    "    keras.layers.Dense(10,name='predictions'),\n",
    "])\n",
    "sequential_model.save_weights('ckpt')\n",
    "load_status = sequential_model.load_weights('ckpt')\n",
    "\n",
    "load_status.assert_consumed()\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"pretrained_model\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "digits (InputLayer)          [(None, 784)]             0         \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 64)                50240     \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 64)                4160      \n",
      "=================================================================\n",
      "Total params: 54,400\n",
      "Trainable params: 54,400\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "inputs = keras.Input(shape=(784,),name='digits')\n",
    "x = keras.layers.Dense(64,activation=\"relu\",name='dense_1')(inputs)\n",
    "x = keras.layers.Dense(64,activation='relu',name='dense_2')(x)\n",
    "outputs = keras.layers.Dense(10,name='predictions')(x)\n",
    "functional_model = keras.Model(inputs=inputs,outputs=outputs,name='3_layer_mlp')\n",
    "\n",
    "pretrained = keras.Model(\n",
    "    functional_model.inputs,functional_model.layers[-1].input,name='pretrained_model'\n",
    ")\n",
    "\n",
    "for w in pretrained.weights:\n",
    "    w.assign(tf.random.normal(w.shape))\n",
    "pretrained.save_weights(\"pretrained_ckpt\")\n",
    "pretrained.summary()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " --------------------------------------------------\n",
      "Model: \"new_model\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "digits (InputLayer)          [(None, 784)]             0         \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 64)                50240     \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 64)                4160      \n",
      "_________________________________________________________________\n",
      "predictions (Dense)          (None, 5)                 325       \n",
      "=================================================================\n",
      "Total params: 54,725\n",
      "Trainable params: 54,725\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "inputs = keras.Input(shape=(784,),name='digits')\n",
    "x = keras.layers.Dense(64,activation='relu',name='dense_1')(inputs)\n",
    "x = keras.layers.Dense(64,activation='relu',name='dense_2')(x)\n",
    "outputs = keras.layers.Dense(5,name='predictions')(x)\n",
    "model = keras.Model(inputs=inputs,outputs=outputs,name='new_model')\n",
    "\n",
    "model.load_weights('pretrained_ckpt')\n",
    "\n",
    "for a,b in zip(pretrained.weights,model.weights):\n",
    "    np.testing.assert_allclose(a.numpy(),b.numpy())\n",
    "print(\"\\n\",\"-\" * 50)\n",
    "model.summary()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "variables:['nested/dense_1/kernel:0', 'nested/dense_1/bias:0', 'nested/dense_2/kernel:0', 'nested/dense_2/bias:0']\n",
      "\n",
      "Changing trainable status of one of the nested layers...\n",
      "\n",
      "variables:['nested/dense_2/kernel:0', 'nested/dense_2/bias:0', 'nested/dense_1/kernel:0', 'nested/dense_1/bias:0']\n",
      "variable ordering changed: True\n"
     ]
    }
   ],
   "source": [
    "class NestedDenseLayer(keras.layers.Layer):\n",
    "    def __init__(self,units,name=None):\n",
    "        super(NestedDenseLayer,self).__init__(name=name)\n",
    "        self.dense_1 = keras.layers.Dense(units,name='dense_1')\n",
    "        self.dense_2 = keras.layers.Dense(units,name='dense_2')\n",
    "\n",
    "    def call(self,inputs):\n",
    "        return self.dense_2(self.dense_1(inputs))\n",
    "\n",
    "nested_model = keras.Sequential([\n",
    "    keras.Input((784,)),\n",
    "    NestedDenseLayer(10,'nested')\n",
    "])\n",
    "variable_names = [v.name for v in nested_model.weights]\n",
    "print(\"variables:{}\".format(variable_names))\n",
    "\n",
    "print(\"\\nChanging trainable status of one of the nested layers...\")\n",
    "nested_model.get_layer('nested').dense_1.trainable = False\n",
    "\n",
    "variable_names_2 =[v.name for v in nested_model.weights]\n",
    "print(\"\\nvariables:{}\".format(variable_names_2))\n",
    "print(\"variable ordering changed:\",variable_names != variable_names_2)\n",
    "\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}